semantic segmentationで道路を青色に塗ってみました
1 はじめに
CX事業本部の平内（SIN）です。
ディープラーニングによる画像認識は、基本的に、Object Detection、Image classification、semantic segmentationの3つとなりますが、今回は、セマンティックセグメンテーションのモデルを使用して、道路を識別する要領を試してみました。
最初に、試してみた様子です。
動画は、Pixels Videosに公開されているものを利用させて頂きました。
2 モデル
今回も、推論に使用したフレームワークは、Open VINO Toolkitです。
そして、モデルは、OpenVINOで利用可能なIR（中間表現フォーマット）へ変換されて公開されている、road-segmentation-adas-0001を利用させて頂きました。
road-segmentation-adas-0001は、各ピクセルをバックグラウンド、道路、縁石、標識の4つのクラスに分類するセグメンテーションモデルです。
入出力は、以下のようになっています。
Inputs
A blob with a BGR image in the format: [B, C=4, H=512, W=896], where:
- B – batch size
- C – number of channels
- H – image height
- W – image width
Outputs
The output is a blob with the shape [B, C=4, H=512, W=896]. It can be treated as a four-channel feature map, where each channel is a probability of one of the classes: BG, road, curb, mark.
3 コード（静止画）
最初に、静止画を処理しているコードです。
RoadSegmentationは、モデルをラップするクラスです。
infer()が、推論を行うメソッドですが、戻り値は、各ピクセルを、4つのクラスに対する信頼度で表現されたTensolです。 しきい値を決めて、一定の値（ここでは、THRESHOLD=0.5としています）となるピクセルを色付けしています。
index.py
import numpy as np
import time
import random
import cv2
import glob
import os
import time
from model import Model
from mrcnn import visualize
from openvino.inference_engine import IEPlugin
class RoadSegmentation(Model):
def __init__(self, plugin, model_path, num_requests=2):
super().__init__(plugin, model_path, num_requests, None)
_, _, h, w = self.input_size
self.__input_height = h
self.__input_width = w
def __prepare_frame(self, frame):
initial_h, initial_w = frame.shape[:2]
scale_h, scale_w = initial_h / float(self.__input_height), initial_w / float(self.__input_width)
in_frame = cv2.resize(frame, (self.__input_width, self.__input_height))
in_frame = in_frame.transpose((2, 0, 1))
in_frame = in_frame.reshape(self.input_size)
return in_frame, scale_h, scale_w
def infer(self, frame):
in_frame, _, _ = self.__prepare_frame(frame)
result = super().infer(in_frame)
# [1, 4, 512, 896] => [4, 512, 896]
return result.squeeze()
# MacOS
device = "CPU"
plugin_dirs = "/opt/intel/openvino/deployment_tools/inference_engine/lib/intel64"
modelPath = "./FP32/"
plugin = IEPlugin(device=device, plugin_dirs = plugin_dirs)
segmentation = RoadSegmentation(plugin, modelPath + "road-segmentation-adas-0001")
input = "./input"
output = "./output"
THRESHOLD= 0.5
WIDTH = 896
HEIGHT = 512
color = (1.0, 0.0, 0.0)
for file in glob.glob("{}/*.png".format(input)):
baseName = os.path.basename(file)
orgImg = cv2.imread(file)
orgImg = cv2.resize(orgImg, (WIDTH, HEIGHT))
segmentImg = orgImg.copy()
start = time.time()
result = segmentation.infer(orgImg)
# result[0]:BG [1]:road [2]:curb [3]:mark
mask = result[1] > THRESHOLD
segmentImg = visualize.apply_mask(segmentImg, mask, color)
print ("processing time:{:.3f} sec {}".format(time.time() - start, baseName))
img = cv2.vconcat([orgImg, segmentImg])
path = "{}/{}".format(output, baseName)
cv2.imwrite(path, img)
下記は、OpenVINOでコンピュータビジョン関連のモデルを使用する場合の、ベースとなるクラスです。
model.py
#
# 下記のコードを参考にさせて頂きました。
# https://github.com/openvinotoolkit/open_model_zoo/blob/master/demos/python_demos/asl_recognition_demo/asl_recognition_demo/common.py
#
from openvino.inference_engine import IENetwork
class Model:
def __init__(self, plugin, model_path, num_requests, output_shape=None):
if model_path.endswith((".xml", ".bin")):
model_path = model_path[:-4]
model = IENetwork(model_path + ".xml", model_path + ".bin")
self.net = plugin.load(network=model)
assert len(self.net.input_info) == 1, "One input is expected"
self.input_name = next(iter(self.net.input_info))
if len(self.net.outputs) > 1:
if output_shape is not None:
candidates = []
for candidate_name in self.net.outputs:
candidate_shape = self.exec_net.requests[0].output_blobs[candidate_name].buffer.shape
if len(candidate_shape) != len(output_shape):
continue
matches = [src == trg or trg < 0
for src, trg in zip(candidate_shape, output_shape)]
if all(matches):
candidates.append(candidate_name)
if len(candidates) != 1:
raise Exception("One output is expected")
self.output_name = candidates[0]
else:
raise Exception("One output is expected")
else:
self.output_name = next(iter(self.net.outputs))
self.input_size = self.net.input_info[self.input_name].input_data.shape
self.output_size = self.net.requests[0].output_blobs[self.output_name].buffer.shape
self.num_requests = num_requests
def infer(self, data):
input_data = {self.input_name: data}
infer_result = self.net.infer(input_data)
return infer_result[self.output_name]
4 出力
上記のコードを実行すると、inputフォルダに置かれた画像を処理し、上下に連結した画像がoutputに出力されます。
そして、コンソールへの出力は、以下のようになっており、各画像ごとの推論及び、着色処理は、0.05sec程度となっていることが分かります。
processing time:0.059 sec img-006.png processing time:0.049 sec img-007.png processing time:0.049 sec img-005.png processing time:0.046 sec img-004.png processing time:0.042 sec img-001.png processing time:0.043 sec img-003.png processing time:0.036 sec img-002.png processing time:0.038 sec img-009.png processing time:0.043 sec img-008.png
推論の出力から、画像に色付けする処理ですが、ピクセルごとループして処理することはもちろん可能です。
しかし、ピクセル単位の画像処理は、非常に負荷が高いため高速化はできません。
今回は、しきい値を超えているかどうかのマスク作成をnumpyの演算、そして、マスクで当該ピクセルに色付けを行う処理は、visualize.apply_mask()を利用しました。
5 コード（動画）
1画像の処理が、0.05sec程度であれば、充分動画にも適用できるということで書いたコードが、以下です。
冒頭の動画は、このコードによるものです。
再生ループの負荷が少しでも下がるように、予め動画のサイズを、モデルの入力サイズに合せました。アスペクト比が維持されていない可能性があるので、ここは、割り切りとさせて下さい。
$ ffmpeg -i input.mp4 -vf scale=896:512 output.mp4
index2.py
import numpy as np
import time
import random
import cv2
import glob
import os
import time
from model import Model
from mrcnn import visualize
from openvino.inference_engine import IEPlugin
class RoadSegmentation(Model):
def __init__(self, plugin, model_path, num_requests=2):
super().__init__(plugin, model_path, num_requests, None)
_, _, h, w = self.input_size
self.__input_height = h
self.__input_width = w
def __prepare_frame(self, frame):
initial_h, initial_w = frame.shape[:2]
scale_h, scale_w = initial_h / float(self.__input_height), initial_w / float(self.__input_width)
in_frame = cv2.resize(frame, (self.__input_width, self.__input_height))
in_frame = in_frame.transpose((2, 0, 1))
in_frame = in_frame.reshape(self.input_size)
return in_frame, scale_h, scale_w
def infer(self, frame):
in_frame, _, _ = self.__prepare_frame(frame)
result = super().infer(in_frame)
# [1, 4, 512, 896] => [4, 512, 896]
return result.squeeze()
# MacOS
device = "CPU"
plugin_dirs = "/opt/intel/openvino/deployment_tools/inference_engine/lib/intel64"
modelPath = "./FP32/"
plugin = IEPlugin(device=device, plugin_dirs = plugin_dirs)
segmentation = RoadSegmentation(plugin, modelPath + "road-segmentation-adas-0001")
VIDEO = "video/output.mp4"
cap = cv2.VideoCapture (VIDEO)
THRESHOLD= 0.5
color = (1.0, 0.0, 0.0)
while(True):
grabbed, frame = cap.read()
if not grabbed: # ループ再生
cap.set(cv2.CAP_PROP_POS_FRAMES, 0)
continue
start = time.time()
result = segmentation.infer(frame)
# result[0]:BG [1]:road [2]:curb [3]:mark
mask = result[1] > THRESHOLD
frame = visualize.apply_mask(frame, mask, color)
print ("processing time:{:.3f} sec".format(time.time() - start))
cv2.imshow('frame', frame)
if cv2.waitKey(1) & 0xFF == ord('q'):
break
cap.release()
cv2.destroyAllWindows()
model.pyは同じです。
6 最後に
今回は、セマンティックセグメンテーションのモデルを使用してみました。モデルのサイズにもよりますが、そこまで処理時間はかかりませんでした。
表示の部分に気をつければ、充分、高速な処理に利用できるのでは、と感じました。
![[Amazon SageMaker] 組み込みアルゴリズム(ObjectDetection)で作成したモデルをOpencVINOで使用してみました](https://devio2023-media.developers.io/wp-content/uploads/2019/05/amazon-sagemaker.png)
